import random
import sqlite3
from datetime import datetime, timedelta

from faker import Faker

# 初始化 Faker 库，用于生成模拟数据
fake = Faker("zh_CN")


def generate_test_data(db_path="data/students.db", count=20000):
    """生成测试数据并插入到SQLite数据库中"""
    try:
        # 连接数据库
        conn = sqlite3.connect(db_path)
        cursor = conn.cursor()

        # 清空现有数据（如果有）
        cursor.execute("DELETE FROM students")

        # 生成并插入测试数据
        for i in range(1, count + 1):
            student_id = f"2023{str(i).zfill(4)}"  # 学号格式：20230001
            name = fake.name()

            # 根据是否团员设置团籍档案状态
            is_tuanyuan_archives_received = random.choice(["是", "否", "非团员"])

            is_student_record_received = random.choice(["是", "否"])

            # 生成最近1年内的随机更新时间
            last_update = (
                datetime.now() - timedelta(days=random.randint(0, 365))
            ).isoformat()

            # 插入数据
            cursor.execute(
                """INSERT INTO students 
                   (id, name, is_tuanyuan_archives_received, is_student_record_received, last_update)
                   VALUES (?, ?, ?, ?, ?)""",
                (
                    student_id,
                    name,
                    is_tuanyuan_archives_received,
                    is_student_record_received,
                    last_update,
                ),
            )

        # 提交事务
        conn.commit()
        print(f"成功生成 {count} 条测试数据")

    except Exception as e:
        print(f"生成测试数据失败: {str(e)}")
        conn.rollback()
    finally:
        conn.close()


if __name__ == "__main__":
    # 生成2000条测试数据
    generate_test_data(count=20000)
